package binarytree

import (
	"fmt"
)

type TreeNode struct {
	data       interface{}
	leftChild  *TreeNode
	rightChild *TreeNode
}

type BinaryTree struct {
	root *TreeNode
}

// 初始化根节点
func NewBinaryTree(v interface{}) *BinaryTree {
	return &BinaryTree{&TreeNode{v, nil, nil}}
}

// 初始化节点
func (this *BinaryTree) NewTreeNode(v interface{}) *TreeNode {
	return &TreeNode{v, nil, nil}
}

// 赋值左子节点
func (this *BinaryTree) SetLeftChild(father *TreeNode, v interface{}) bool {
	// 如果左子节点不为空，赋值失败
	if father.leftChild != nil {
		return false
	}

	// 新建节点赋值给左孩子节点
	father.leftChild = this.NewTreeNode(v)
	return true
}

// 赋值右子节点
func (this *BinaryTree) SetRightChild(father *TreeNode, v interface{}) bool {
	// 如果右子节点不为空，赋值失败
	if father.rightChild != nil {
		return false
	}

	// 新建节点赋值给右孩子节点
	father.rightChild = this.NewTreeNode(v)
	return true
}

// 获取根节点
func (this *BinaryTree) GetRoot() *TreeNode {
	return this.root
}

// 获取左子节点
func (this *BinaryTree) GetLeftChild(father *TreeNode) *TreeNode {
	return father.leftChild
}

// 获取右子节点
func (this *BinaryTree) GetRightChild(father *TreeNode) *TreeNode {
	return father.rightChild
}

// 前序遍历二叉树
// func (this *BinaryTree) PreOrderTraversal() []string {

// 	var stack []*TreeNode
// 	res := []string{}

// 	// 根节点为空
// 	if this.root == nil {
// 		return nil
// 	}

// 	// 根节点的左右节点为空
// 	if this.root.leftChild == nil && this.root.rightChild == nil {
// 		return []string{this.root.data.(string)}
// 	}

// 	stack = append(stack, this.root)
// 	for len(stack) != 0 {
// 		node := stack[len(stack)-1]
// 		stack = stack[:len(stack)-1]
// 		res = append(res, node.data.(string))

// 		if node.rightChild != nil {
// 			stack = append(stack, node.rightChild)
// 		}

// 		if node.leftChild != nil {
// 			stack = append(stack, node.leftChild)
// 		}

// 	}

// 	return res
// }

// 前序遍历
func (this *BinaryTree) PreOrderTraversal(node *TreeNode) []string {

	// 初始化
	res, leftRes, rightRes := []string{}, []string{}, []string{}

	// 如果父节点不存在，返回空
	if node == nil {
		return []string{}
	}

	// 获取左子节点
	leftRes = this.PreOrderTraversal(node.leftChild)
	// 获取右子节点
	rightRes = this.PreOrderTraversal(node.rightChild)

	res = append(res, node.data.(string))
	res = append(res, leftRes...)
	res = append(res, rightRes...)

	return res
}

// 中序遍历
func (this *BinaryTree) InOrderTraversal(node *TreeNode) []string {

	res, leftRes, rightRes := []string{}, []string{}, []string{}

	if node == nil {
		return res
	}

	// 获取左子节点
	leftRes = this.InOrderTraversal(node.leftChild)
	// 获取右子节点
	rightRes = this.InOrderTraversal(node.rightChild)

	res = append(res, leftRes...)
	res = append(res, node.data.(string))
	res = append(res, rightRes...)

	return res
}

// 后序遍历
func (this *BinaryTree) PostOrderTraversal(node *TreeNode) []string {

	res, leftRes, rightRes := []string{}, []string{}, []string{}

	if node == nil {
		return res
	}

	// 获取左子节点
	leftRes = this.PostOrderTraversal(node.leftChild)
	// 获取右子节点
	rightRes = this.PostOrderTraversal(node.rightChild)

	res = append(res, leftRes...)
	res = append(res, rightRes...)
	res = append(res, node.data.(string))

	return res
}

// 打印二叉树的节点（前序|中序|后序）
func (this *BinaryTree) PrintTree(status string) {

	res := []string{}
	switch status {
	case "pre":
		res = this.PreOrderTraversal(this.root)
	case "in":
		res = this.InOrderTraversal(this.root)
	case "post":
		res = this.PostOrderTraversal(this.root)
	default:
		return
	}

	// 打印
	for i := 0; i < len(res); i++ {
		fmt.Printf("%v", res[i])

		if i != len(res)-1 {
			fmt.Print("-->")
		}
	}

	fmt.Println()
}
